#include "viml/liveness_analyzer.h"
#include "viml/instruction.h"
#include "viml/value.h"
#include <algorithm>
#include <iostream>

namespace viml {

bool LivenessAnalyzer::analyze(const std::shared_ptr<Bytecode>& bytecode) {
    if (!bytecode) {
        return false;
    }
    
    reset();
    
    instruction_count_ = bytecode->instruction_count();
    if (instruction_count_ == 0) {
        return true;
    }
    
    // 初始化活性信息数组
    liveness_info_.resize(instruction_count_);
    
    // 保存指令副本用于控制流分析
    instructions_.clear();
    instructions_.reserve(instruction_count_);
    for (size_t i = 0; i < instruction_count_; ++i) {
        instructions_.push_back(bytecode->get_instruction(i));
    }
    
    // 阶段1: 计算每条指令的def和use集合
    compute_def_use_sets(bytecode);
    
    // 阶段2: 迭代计算live_in和live_out
    compute_liveness();
    
    // 阶段3: 构建寄存器生命周期区间
    build_live_intervals();
    
    // 阶段4: 构建寄存器冲突图
    build_conflict_graph();
    
    return true;
}

const LivenessInfo& LivenessAnalyzer::get_liveness_at(size_t pc) const {
    if (pc >= liveness_info_.size()) {
        static LivenessInfo empty;
        return empty;
    }
    return liveness_info_[pc];
}

void LivenessAnalyzer::reset() {
    liveness_info_.clear();
    live_intervals_.clear();
    conflict_graph_.clear();
    instructions_.clear();
    instruction_count_ = 0;
    max_live_registers_ = 0;
}

void LivenessAnalyzer::compute_def_use_sets(const std::shared_ptr<Bytecode>& bytecode) {
    // TODO: 实现def/use集合计算
    // 遍历每条指令,分析其定义和使用的寄存器
    
    for (size_t pc = 0; pc < instruction_count_; ++pc) {
        const auto& instr = bytecode->get_instruction(pc);
        analyze_instruction_def_use(instr, pc);
    }
}

void LivenessAnalyzer::analyze_instruction_def_use(const Instruction& instr, size_t pc) {
    auto& info = liveness_info_[pc];
    OpCode op = instr.opcode();
    
    switch (op) {
        // 数据移动指令 - def A, use 无/B
        case OpCode::LOAD_CONST:
        case OpCode::LOAD_NIL:
        case OpCode::LOAD_TRUE:
        case OpCode::LOAD_FALSE:
            info.def_set.insert(instr.a());
            break;
            
        case OpCode::MOVE:  // MOVE A, B -> def A, use B
            info.def_set.insert(instr.a());
            info.use_set.insert(instr.b());
            break;
        
        // 全局变量操作
        case OpCode::LOAD_GLOBAL:  // LOAD_GLOBAL A, Bx -> def A
        case OpCode::LOAD_LOCAL:   // LOAD_LOCAL A, B -> def A, use B(作为索引)
            info.def_set.insert(instr.a());
            break;
            
        case OpCode::STORE_GLOBAL:  // STORE_GLOBAL A, Bx -> use A
        case OpCode::STORE_LOCAL:   // STORE_LOCAL A, B -> use A
            info.use_set.insert(instr.a());
            break;
        
        // 作用域变量操作
        case OpCode::LOAD_SCOPE:  // LOAD_SCOPE A, ... -> def A
            info.def_set.insert(instr.a());
            break;
            
        case OpCode::STORE_SCOPE:  // STORE_SCOPE ..., B -> use B
            info.use_set.insert(instr.a());  // a字段存储value_reg
            break;
        
        // 算术运算指令 - def A, use B C
        case OpCode::ADD:
        case OpCode::SUB:
        case OpCode::MUL:
        case OpCode::DIV:
        case OpCode::MOD:
            info.def_set.insert(instr.a());
            info.use_set.insert(instr.b());
            info.use_set.insert(instr.c());
            break;
            
        case OpCode::NEG:  // NEG A, B -> def A, use B
            info.def_set.insert(instr.a());
            info.use_set.insert(instr.b());
            break;
        
        // 比较运算指令 - def A, use B C
        case OpCode::EQ:
        case OpCode::NE:
        case OpCode::LT:
        case OpCode::LE:
        case OpCode::GT:
        case OpCode::GE:
        case OpCode::AND:
        case OpCode::OR:
            info.def_set.insert(instr.a());
            info.use_set.insert(instr.b());
            info.use_set.insert(instr.c());
            break;
            
        case OpCode::NOT:  // NOT A, B -> def A, use B
            info.def_set.insert(instr.a());
            info.use_set.insert(instr.b());
            break;
        
        // 控制流指令
        case OpCode::JMP:  // JMP Bx -> 无def/use
            break;
            
        case OpCode::JMP_IF_TRUE:   // JMP_IF_TRUE A, Bx -> use A
        case OpCode::JMP_IF_FALSE:  // JMP_IF_FALSE A, Bx -> use A
            info.use_set.insert(instr.a());
            break;
        
        // 迭代指令
        case OpCode::ITER_INIT:  // ITER_INIT A, B -> def A, A+1, use B
            info.def_set.insert(instr.a());
            info.def_set.insert(instr.a() + 1);
            info.use_set.insert(instr.b());
            break;
            
        case OpCode::ITER_NEXT:  // ITER_NEXT A, Bx -> def A+2, use A
            info.def_set.insert(instr.a() + 2);
            info.use_set.insert(instr.a());
            break;
        
        // 函数调用指令
        case OpCode::CALL: {  // CALL A, B, C -> def A..A+C-1, use B, A+1..A+B
            uint8_t result = instr.a();
            uint8_t func = instr.b();
            uint8_t argc = instr.c();
            
            // def: result寄存器
            info.def_set.insert(result);
            
            // use: 函数寄存器和参数寄存器
            info.use_set.insert(func);
            for (uint8_t i = 1; i <= argc; ++i) {
                info.use_set.insert(result + i);
            }
            break;
        }
        
        case OpCode::RETURN_OP: {  // RETURN A, B -> use A..A+B-1
            uint8_t start = instr.a();
            uint8_t count = instr.b();
            for (uint8_t i = 0; i < count; ++i) {
                info.use_set.insert(start + i);
            }
            break;
        }
        
        case OpCode::FUNC_DEF:  // FUNC_DEF A, Bx -> def A
            info.def_set.insert(instr.a());
            break;
        
        // 数据结构操作
        case OpCode::NEW_LIST:  // NEW_LIST A, B -> def A
        case OpCode::NEW_DICT:  // NEW_DICT A -> def A
            info.def_set.insert(instr.a());
            break;
            
        case OpCode::GET_INDEX:  // GET_INDEX A, B, C -> def A, use B C
            info.def_set.insert(instr.a());
            info.use_set.insert(instr.b());
            info.use_set.insert(instr.c());
            break;
            
        case OpCode::SET_INDEX:  // SET_INDEX A, B, C -> use A B C
            info.use_set.insert(instr.a());
            info.use_set.insert(instr.b());
            info.use_set.insert(instr.c());
            break;
            
        case OpCode::LEN:  // LEN A, B -> def A, use B
            info.def_set.insert(instr.a());
            info.use_set.insert(instr.b());
            break;
            
        case OpCode::LIST_APPEND:  // LIST_APPEND A, B -> use A B
            info.use_set.insert(instr.a());
            info.use_set.insert(instr.b());
            break;
            
        case OpCode::DICT_SET:  // DICT_SET A, B, C -> use A B C
            info.use_set.insert(instr.a());
            info.use_set.insert(instr.b());
            info.use_set.insert(instr.c());
            break;
        
        // 字符串操作
        case OpCode::STR_CAT:  // STR_CAT A, B, C -> def A, use B C
            info.def_set.insert(instr.a());
            info.use_set.insert(instr.b());
            info.use_set.insert(instr.c());
            break;
        
        // 调试指令
        case OpCode::BREAKPOINT:
        case OpCode::DBG_LINE:
            break;
        
        // 内置函数
        case OpCode::ECHO: {  // ECHO A, B -> use A..A+B-1
            uint8_t start = instr.a();
            uint8_t count = instr.b();
            for (uint8_t i = 0; i < count; ++i) {
                info.use_set.insert(start + i);
            }
            break;
        }
        
        // 控制指令
        case OpCode::HALT:
            break;
    }
}

void LivenessAnalyzer::compute_liveness() {
    // TODO: 实现活性分析迭代算法
    //
    // 算法:
    // 1. 初始化所有live_in和live_out为空
    // 2. 重复以下过程直到不动点:
    //    for each instruction i in reverse order:
    //        live_out[i] = union of live_in[succ] for all successors succ of i
    //        live_in[i] = use[i] ∪ (live_out[i] - def[i])
    // 3. 收敛条件: 所有live_in和live_out不再变化
    
    bool changed = true;
    int iteration = 0;
    const int max_iterations = 100;  // 防止无限循环
    
    while (changed && iteration < max_iterations) {
        changed = false;
        iteration++;
        
        // 反向遍历指令
        for (size_t i = instruction_count_; i > 0; --i) {
            size_t pc = i - 1;
            auto& info = liveness_info_[pc];
            
            // 保存旧的live_in用于检测变化
            auto old_live_in = info.live_in;
            
            // 计算live_out: 所有后继指令的live_in的并集
            std::unordered_set<uint32_t> new_live_out;
            for (size_t succ_pc : get_successors(pc, instructions_[pc])) {
                const auto& succ_info = liveness_info_[succ_pc];
                new_live_out.insert(succ_info.live_in.begin(), succ_info.live_in.end());
            }
            
            // 计算live_in: use ∪ (live_out - def)
            std::unordered_set<uint32_t> new_live_in = info.use_set;
            for (uint32_t reg : new_live_out) {
                if (info.def_set.find(reg) == info.def_set.end()) {
                    new_live_in.insert(reg);
                }
            }
            
            // 更新并检测变化
            info.live_out = std::move(new_live_out);
            if (new_live_in != old_live_in) {
                info.live_in = std::move(new_live_in);
                changed = true;
            }
            
            // 更新峰值寄存器数
            max_live_registers_ = std::max(max_live_registers_, info.live_out.size());
        }
    }
    
    if (iteration >= max_iterations) {
        std::cerr << "Warning: Liveness analysis did not converge" << std::endl;
    }
}

std::vector<size_t> LivenessAnalyzer::get_successors(size_t pc, const Instruction& instr) const {
    std::vector<size_t> successors;
    
    if (pc >= instruction_count_) {
        return successors;
    }
    
    OpCode op = instr.opcode();
    
    // 根据指令类型确定后继指令
    switch (op) {
        case OpCode::JMP: {
            // 无条件跳转: 只有跳转目标
            int16_t offset = instr.sbx();
            size_t target = pc + 1 + offset;
            if (target < instruction_count_) {
                successors.push_back(target);
            }
            break;
        }
        
        case OpCode::JMP_IF_TRUE:
        case OpCode::JMP_IF_FALSE: {
            // 条件跳转: 两个后继(跳转和顺序执行)
            int16_t offset = instr.sbx();
            size_t target = pc + 1 + offset;
            if (target < instruction_count_) {
                successors.push_back(target);
            }
            if (pc + 1 < instruction_count_) {
                successors.push_back(pc + 1);
            }
            break;
        }
        
        case OpCode::ITER_NEXT: {
            // 迭代跳转: 两个后继(跳转和顺序执行)
            int16_t offset = instr.sbx();
            size_t target = pc + 1 + offset;
            if (target < instruction_count_) {
                successors.push_back(target);
            }
            if (pc + 1 < instruction_count_) {
                successors.push_back(pc + 1);
            }
            break;
        }
        
        case OpCode::RETURN_OP:
        case OpCode::HALT:
            // 返回和停止: 无后继
            break;
        
        default:
            // 普通指令: 顺序执行
            if (pc + 1 < instruction_count_) {
                successors.push_back(pc + 1);
            }
            break;
    }
    
    return successors;
}

void LivenessAnalyzer::build_live_intervals() {
    // TODO: 从活性信息构建生命周期区间
    //
    // 算法:
    // 1. 对每个寄存器r:
    //    - 找到r首次在def_set中出现的位置(定义点)
    //    - 找到r最后在use_set或live_out中出现的位置(最后使用点)
    //    - 记录[定义点, 最后使用点]为r的生命周期
    // 2. 如果寄存器有多次定义,为每次定义创建独立的区间
    
    std::unordered_map<uint32_t, size_t> first_def;
    std::unordered_map<uint32_t, size_t> last_use;
    
    // 遍历所有指令,记录每个寄存器的首次定义和最后使用
    for (size_t pc = 0; pc < instruction_count_; ++pc) {
        const auto& info = liveness_info_[pc];
        
        // 记录定义点
        for (uint32_t reg : info.def_set) {
            if (first_def.find(reg) == first_def.end()) {
                first_def[reg] = pc;
            }
        }
        
        // 记录使用点
        for (uint32_t reg : info.use_set) {
            last_use[reg] = pc;
        }
        
        // 检查live_out中的寄存器
        for (uint32_t reg : info.live_out) {
            last_use[reg] = pc;
        }
    }
    
    // 构建生命周期区间
    for (const auto& [reg, start] : first_def) {
        size_t end = last_use.count(reg) ? last_use[reg] : start;
        live_intervals_.emplace_back(reg, start, end);
    }
}

void LivenessAnalyzer::build_conflict_graph() {
    // TODO: 从生命周期区间构建冲突图
    //
    // 算法:
    // 1. 对每对寄存器(r1, r2):
    //    - 如果它们的生命周期区间有重叠,则在冲突图中添加边(r1, r2)
    // 2. 冲突图用于后续的寄存器分配(图着色算法)
    
    for (size_t i = 0; i < live_intervals_.size(); ++i) {
        const auto& interval1 = live_intervals_[i];
        uint32_t reg1 = interval1.register_id;
        
        for (size_t j = i + 1; j < live_intervals_.size(); ++j) {
            const auto& interval2 = live_intervals_[j];
            uint32_t reg2 = interval2.register_id;
            
            // 检查生命周期是否冲突
            if (interval1.conflicts_with(interval2)) {
                conflict_graph_[reg1].insert(reg2);
                conflict_graph_[reg2].insert(reg1);
            }
        }
    }
}

} // namespace viml
